#
# messkit.async_stream
#
# Copyright 2006 Helsinki Institute for Information Technology (HIIT)
# and the authors.
#
# Authors: Ken Rimey <rimey@hiit.fi>
#

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""
Stream implementation based on asyncore

This is based in the usual way on a event loop shared by all streams
and servers, but it takes the unusual approach of running the event
loop in a helper thread.  This implementation integrates nicely with
tlslite (see tls_async_stream.py).
"""

from __future__ import generators

import os
import traceback
import socket

# We need the Python 2.4 version of asyncore for its support for
# private socket maps.
from messkit.asyncore import dispatcher, file_dispatcher, loop
try:
    from messkit import asyncore
    del asyncore.socket_map   # Don't accidentally use the global map.
except AttributeError:
    pass

from messkit.async import Thread
from messkit.stream import Stream as _Stream
from messkit.stream import StreamServer as _StreamServer
from messkit.stream import StreamManager as _StreamManager

__all__ = ['Stream', 'StreamServer', 'StreamManager']

class Channel(dispatcher):
    def __init__(self, map):
        dispatcher.__init__(self, map=map)

        # We have two queues of handlers.  Each handler is a
        # generator, which will yield True as long as it wants to
        # remain in effect.  This convention, where the handler can
        # either yield False or raise StopIteration to signal that it
        # is finished, makes it easier to write handlers that always
        # finish their work in the first call to next().
        self.read_handlers = []
        self.write_handlers = []

    def stop(self):
        self.del_channel()

        # Break cycles and immediately disable the callback mechanism.
        self.read_handlers = ()
        self.write_handlers = ()

    def close(self):
        self.stop()
        self.socket.close()

    def readable(self):
        return self.read_handlers

    def writable(self):
        return self.write_handlers

    def handle_read(self):
        handlers = self.read_handlers
        if handlers:
            generator = handlers.pop(0)
            for more in generator:
                if more:
                    handlers.insert(0, generator)
                break

    def handle_write(self):
        handlers = self.write_handlers
        if handlers:
            generator = handlers.pop(0)
            for more in generator:
                if more:
                    handlers.insert(0, generator)
                break

    def handle_connect(self):
        # Note that this is sometimes called directly from connect(),
        # in the main thread.
        pass

    def handle_accept(self):
        self.handle_read()

    def handle_close(self):
        pass                    # Don't close channel here.

    def handle_error(self):
        traceback.print_exc()   # I find asyncore's tracebacks unreadable.

class Stream(_Stream):
    def _init(self, sock, aggressive_write=True, channel_factory=Channel):
        self.aggressive_write = aggressive_write

        self.channel = channel_factory(self.mgr.map)
        if isinstance(sock, tuple):
            self.channel.create_socket(*sock)
        elif isinstance(sock, socket.SocketType):
            sock.setblocking(0)
            self.channel.connected = True
            self.channel.set_socket(sock)
        else:
            raise TypeError

        self.channel.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)

    def _close(self, cb):
        generator = self.__handle_close(cb)
        if not self.connected:
            self.mgr.async(generator.next)
        else:
            self.schedule_write_handler(generator)

    def __handle_close(self, cb):
        self.channel.close()
        self.issue_callback(cb)
        yield False

    def _connect(self, address, cb):
        try:
            self.channel.connect(address)
        except socket.error, e: # E.g., error 64.
            self.issue_callback(cb, e)
        else:
            self.schedule_write_handler(self.__handle_connect(cb))

    def __handle_connect(self, cb):
        errno = self.channel.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
        if errno == 0:
            error = None
        else:
            error = socket.error(errno, os.strerror(errno))
            self.error = error

        self.issue_callback(cb, error)
        yield False

    def _flush(self, buf):
        if not self.channel.write_handlers:
            self.schedule_write_handler(self.__handle_flush(buf))

    def __handle_flush(self, buf):
        while buf.size and not self.error:
            s = buf.peek(self.send_size)
            try:
                n = self.channel.send(s)
            except socket.error, e:
                self._signal_error(e)
                return

            if n:
                self.bytes_out += n
                buf.skip(n)

            if n < len(s) or (not self.aggressive_write and buf.size):
                yield True

    def _fill(self, buf, size, cb):
        start = self.bytes_in
        stop = start + size
        self.schedule_read_handler(self.__handle_fill(buf, start, stop, cb))

    def __handle_fill(self, buf, start, stop, cb):
        while self.bytes_in < stop and not self.eof_seen:
            try:
                s = self.channel.recv(self.recv_size)
            except socket.error, e:
                self._signal_error(e)
                return

            if not s:
                self.eof_seen = True
            else:
                self.bytes_in += len(s)
                buf.append(s)

            if self.bytes_in < stop and not self.eof_seen:
                yield True

        self.issue_callback(cb, self.bytes_in - start)

    def _get_socket(self):
        return self.channel.socket

    def schedule_read_handler(self, generator):
        self.channel.read_handlers.append(generator)
        self.mgr.notify()

    def schedule_write_handler(self, generator):
        self.channel.write_handlers.append(generator)
        self.mgr.notify()

class StreamServer(_StreamServer):
    def _init(self, params):
        self.channel = Channel(self.mgr.map)
        self.channel.create_socket(*params)
        self.channel.set_reuse_addr()

    def _close(self, cb):
        self.mgr.async(self.__close, cb)

    def __close(self, cb):
        self.channel.close()
        self.issue_callback(cb)

    def _listen(self, address):
        self.channel.bind(address)
        self.channel.listen(5)

    def _accept(self, cb):
        self.channel.read_handlers.append(self.__handle_accept(cb))
        self.mgr.notify()

    def __handle_accept(self, cb):
        result = self.channel.accept()

        while result is None:   # Probably never happens.
            yield True
            result = self.channel.accept()

        self.issue_callback(cb, *result)
        yield False

class ControlChannel(file_dispatcher):
    def __init__(self, map):
        self.queue = []
        r, w = os.pipe()
        self.__w = w
        file_dispatcher.__init__(self, r, map)

    def async(self, f, *args):
        self.queue.append((f, args))
        self.notify()

    def notify(self):
        os.write(self.__w, '!')

    def terminate(self):
        os.close(self.__w)

    def handle_read(self):
        self.recv(1000)
        while self.queue:
            f, args = self.queue.pop(0)
            f(*args)

    def handle_close(self):
        self.close()

    def writable(self):
        return False        # I have seen write events under Mac OS X.

class StreamManager(_StreamManager):
    def __init__(self, callback_callback=None):
        _StreamManager.__init__(self, callback_callback)

        self.map = {}
        self.cc = ControlChannel(self.map)
        self.async = self.cc.async
        self.notify = self.cc.notify

        self.thread = Thread()
        self.thread.async(lambda: loop(map=self.map, timeout=10.0))

    def _create_stream(self, *args):
        return Stream(self, *args)

    def _create_server(self, params):
        return StreamServer(self, params)

    def _close(self, cb):
        self.cc.terminate()
        self.thread.async(lambda: self.issue_main_thread_callback(cb))
        self.thread.terminate(timeout=None)
